import os
import cv2
import numpy as np
from modelscope.msdatasets import MsDataset

def create_directory_structure():
    """创建多语言数据集的目录结构"""
    # 创建基础数据目录
    base_dir = 'handwriting_data'
    if not os.path.exists(base_dir):
        os.makedirs(base_dir)
    
    # 创建语言类别目录
    lang_dirs = {
        'digit': list(range(10)),  # 数字0-9
        'english': [chr(i) for i in range(65, 91)],  # 大写字母A-Z
        'chinese': []  # 中文字符（后续手动添加）
    }
    
    for lang, chars in lang_dirs.items():
        lang_path = os.path.join(base_dir, lang)
        if not os.path.exists(lang_path):
            os.makedirs(lang_path)
        
        # 为每个字符创建子目录
        for char in chars:
            char_path = os.path.join(lang_path, str(char))
            if not os.path.exists(char_path):
                os.makedirs(char_path)
    
    return base_dir

def prepare_mnist_data():
    """准备MNIST数据集"""
    # 使用ModelScope加载MNIST数据集
    train_ds = MsDataset.load('modelscope/mnist', subset_name='mnist', split='train')
    test_ds = MsDataset.load('modelscope/mnist', subset_name='mnist', split='test')
    
    base_dir = create_directory_structure()
    digit_dir = os.path.join(base_dir, 'digit')
    
    # 处理训练集和测试集
    for idx, dataset in enumerate([train_ds, test_ds]):
        prefix = 'train' if idx == 0 else 'test'
        
        for i, item in enumerate(dataset):
            if i % 1000 == 0:
                print(f"处理第 {i} 张{prefix}图片...")
            
            # 将PIL Image转换为numpy数组
            image = np.array(item['image'])
            label = item['label']
            
            # 保存图片到对应的数字目录
            image_path = os.path.join(digit_dir, str(label), f'mnist_{prefix}_{i}.png')
            cv2.imwrite(image_path, image)

def add_custom_image(image_path, lang_type, char_label):
    """
    添加自定义手写体图片
    
    参数:
        image_path (str): 图片路径
        lang_type (str): 语言类型 ('digit', 'english', 'chinese')
        char_label (str): 字符标签
    """
    base_dir = 'handwriting_data'
    target_dir = os.path.join(base_dir, lang_type, str(char_label))
    
    # 如果目标目录不存在，创建它
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    
    # 生成目标文件名
    filename = os.path.basename(image_path)
    target_path = os.path.join(target_dir, f'custom_{filename}')
    
    # 复制并预处理图片
    image = cv2.imread(image_path)
    if image is not None:
        # 转换为灰度图
        if len(image.shape) == 3:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 调整大小为28x28
        image = cv2.resize(image, (28, 28))
        
        # 保存处理后的图片
        cv2.imwrite(target_path, image)
        print(f"成功添加图片到 {target_path}")
    else:
        print(f"无法读取图片: {image_path}")

def main():
    # 准备基础的MNIST数据集
    prepare_mnist_data()
    
    # 示例：如何添加自定义图片
    # add_custom_image('path/to/your/image.jpg', 'chinese', '你')
    # add_custom_image('path/to/your/image.jpg', 'english', 'A')

if __name__ == '__main__':
    main() 